fix(mcp): MCP Instrumentation: streamablehttp_client Parameter Corruption #3199
Conversation
|
|
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughThe MCP instrumentor's Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant McpInstrumentor
participant WrappedTransport
Caller ->> McpInstrumentor: call traced_method()
McpInstrumentor ->> WrappedTransport: invoke wrapped transport
WrappedTransport -->> McpInstrumentor: returns result (2-tuple or 3-tuple or other)
alt result is 2-tuple
McpInstrumentor ->> McpInstrumentor: wrap reader/writer
McpInstrumentor -->> Caller: yield wrapped (reader, writer)
else result is 3-tuple
McpInstrumentor ->> McpInstrumentor: wrap reader/writer + include callback
McpInstrumentor -->> Caller: yield (reader, writer, callback)
else unpack/error
McpInstrumentor ->> McpInstrumentor: log warning/error
McpInstrumentor -->> Caller: yield original result
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15 minutes Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![ellipsis-dev[bot]](https://avatars.githubusercontent.com/in/64358?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Changes requested ❌
Reviewed everything up to baef8ef in 1 minute and 56 seconds. Click for details.
- Reviewed
30lines of code in1files - Skipped
0files when reviewing. - Skipped posting
1draft comments. View those below. - Modify your settings and rules to customize what types of comments Ellipsis leaves. And don't forget to react with 👍 or 👎 to teach Ellipsis.
1. packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py:132
- Draft comment:
Using try/except to distinguish a 2-tuple from a 3-tuple is brittle. Consider checking the length of “result” (e.g. using if len(result)==3) to decide which unpacking to use, and update the async generator’s type hint accordingly. - Reason this comment was not posted:
Decided after close inspection that this draft comment was likely wrong and/or not actionable: usefulness confidence = 20% vs. threshold = 50% The comment makes a valid point about using try/except for flow control being somewhat brittle. Using len() would be more explicit and clearer. However, the code was just changed to add better error handling, not to change the core unpacking logic. The suggestion is more of a code quality improvement than a critical issue. The try/except approach could actually be more robust in some cases, as it handles duck typing better than checking lengths. The performance difference is likely negligible. While the len() approach might be marginally clearer, the current approach works and the code changes were focused on error handling, not refactoring the tuple unpacking logic. This comment suggests a valid but non-critical code quality improvement that is tangential to the actual changes being made. Per the rules, we should not keep comments that are purely informative or suggest non-critical refactors.
Workflow ID: wflow_dRm5wie2RSW0OlsL
You can customize 
packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py
Show resolved
Hide resolved
feng-95
changed the title
|
@feng-95 can you look into the comment from Elipsis? |
Since I'm unsure whether creating a new wrapper function specifically for |
|
Hi @nirga , could you please take a look to this again? Would you think it will be better to creating a new wrapper function specifically for |
packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py
Show resolved
Hide resolved
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py (1)
126-132: Consider standardizing the return type for better API consistency.While the Union type correctly represents the two possible return types (2-tuple and 3-tuple), having inconsistent return types can make the API harder to use for consumers. They need to handle both cases, which increases complexity.
Consider one of these approaches:
- Always return a 3-tuple, using
Nonefor the third parameter when not present- Create separate wrapper methods for different transport types
Example of approach 1:
- ) -> AsyncGenerator[ - Union[ - Tuple[InstrumentedStreamReader, InstrumentedStreamWriter], - Tuple[InstrumentedStreamReader, InstrumentedStreamWriter, Any] - ], - None + ) -> AsyncGenerator[ + Tuple[InstrumentedStreamReader, InstrumentedStreamWriter, Any], + None
🧹 Nitpick comments (2)
packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py (2)
139-147: Consider more specific exception handling in the nested try block.The nested try-except block catches all exceptions, which might mask other issues beyond tuple unpacking. Consider catching
ValueErrorspecifically for the 3-tuple unpacking attempt.except ValueError: - try: + try: read_stream, write_stream, get_session_id_callback = result yield InstrumentedStreamReader( read_stream, tracer ), InstrumentedStreamWriter(write_stream, tracer), get_session_id_callback - except Exception as e: + except ValueError as e: logging.warning(f"mcp instrumentation _transport_wrapper exception: {e}") yield result + except Exception as e: + logging.warning(f"mcp instrumentation _transport_wrapper unexpected exception: {e}") + yield result
146-147: Improve warning messages for better debugging.The warning messages could be more descriptive to help identify which specific unpacking attempt failed and what the actual result structure was.
- logging.warning(f"mcp instrumentation _transport_wrapper exception: {e}") + logging.warning(f"MCP instrumentation: Failed to unpack transport result as 3-tuple: {e}, result type: {type(result)}") yield result except Exception as e: - logging.warning(f"mcp instrumentation transport_wrapper exception: {e}") + logging.warning(f"MCP instrumentation: Unexpected error in transport wrapper: {e}, result type: {type(result)}")Also applies to: 149-150
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py(2 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py (1)
packages/opentelemetry-instrumentation-mcp/tests/conftest.py (1)
tracer(48-49)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Lint
- GitHub Check: Build Packages (3.11)
🔇 Additional comments (2)
packages/opentelemetry-instrumentation-mcp/opentelemetry/instrumentation/mcp/instrumentation.py (2)
3-3: LGTM! The Union import is correctly added.The import of
Unionis necessary to support the updated return type annotation that can yield either 2-tuple or 3-tuple results.
133-150: The fix correctly handles both 2-tuple and 3-tuple returns, addressing the original issue.The implementation properly:
- First attempts to unpack as a 2-tuple for
stdio_clientandsse_client- Falls back to 3-tuple unpacking for
streamablehttp_clientwith the session ID callback- Includes comprehensive error handling with logging for debugging
This resolves the
ValueErrorthat was occurring during startup.
feat(instrumentation): ...orfix(instrumentation): ....When using the MCP Instrumentation to instrument the mcp streamablehttp_client, the parameters returned by the instrumentation are incorrect. The streamablehttp_client returns
https://github.com/modelcontextprotocol/python-sdk/blob/v1.9.0/src/mcp/client/streamable_http.py#L486
But the instrumentated function only returns
read_stream, write_stream.Important
Fixes parameter return issue in MCP Instrumentation's
streamablehttp_clientby correctly handlingget_session_id_callbackininstrumentation.py.traced_methodwithin_transport_wrapperininstrumentation.py.get_session_id_callbackwhen present, aligning withstreamablehttp_client's expected return values._transport_wrapperto aid debugging.This description was created by
for baef8ef. You can customize this summary. It will automatically update as commits are pushed.
Summary by CodeRabbit
Bug Fixes
Chores